close all
clear

folderName = '1_Bloch_oscillations_freq1_780Hz';
prefix = 'scan';
paramVar ='evolution_dur';

col1 = 115; 
col2 = 128;
row1 = 94; %difference must be even
row2 = 106; %starting site must be in the center
centersite = (row2-row1)/2+1;

[param_names, param_table] = get_batch_params(folderName); %param_table has extra zero at the end for some reason
paramInd = find(strcmp(param_names, paramVar));
alltvalues = param_table(:, paramInd);
alltlength = length(alltvalues);
atomMatrix = zeros(248,248);
tlen = numel(unique(alltvalues));


%%

V0List = 780;
JFit = zeros(length(V0List),4);

for vv=1:length(V0List)
    V0 = V0List(vv);
    start = 1+tlen*(vv-1); stop = tlen*vv;%tlength;
    tvalues = alltvalues(start:stop);
    Nrealizations = zeros(stop-start+1,1);
    Npost = ones(stop-start+1,1);
    bo = zeros(stop-start+1,row2-row1+1);
    
    for jj = 1:(stop-start+1) %each time step
        listing = dir(fullfile(folderName, [prefix '*' num2str(start+jj-1,'%03.f') 'atomMatrix.mat'])); %find the file
        Nshots = size(listing,1); %count how many shots taken by count how many files having this title
        Nrealizations(jj) = (col2-col1+1) * Nshots; % = how many runs each shot * howm many shots in total
        for i = 1:Nshots
            load(fullfile(folderName, listing(i).name)); % load file one by one
            for ii = col1:col2 %for each qunatum walker
                if sum(atomMatrix(row1:row2,ii)) == 1 %see if the atom is there
                    Npost(jj) = Npost(jj) + 1; % count if this shot is useful
                    bo(jj,:) = bo(jj,:) + atomMatrix(row1:row2,ii)'; %count atom at each site
                end 
            end 
        end
    end
     
    bo = (bo./repmat(Npost,1,size(bo,2)))';
  
    bovar = zeros(stop-start+1,row2-row1+1);
    for jj = 1:(stop-start+1) %each time step
        listing = dir(fullfile(folderName, [prefix '*' num2str(start+jj-1,'%03.f') 'atomMatrix.mat'])); %find the file
        Nshots = size(listing,1); %count how many shots taken by count how many files having this title
        for i = 1:Nshots
            load(fullfile(folderName, listing(i).name)); % load file one by one
            for ii = col1:col2 %for each qunatum walker
                if sum(atomMatrix(row1:row2,ii)) == 1 %see if the atom is there
                    bostar = bo';
                    bovar(jj,:) = bovar(jj,:) +   (atomMatrix(row1:row2,ii)'-bostar(jj,:)).^2; %count atom at each site
                end 
            end 
        end
    end
    bovar = (bovar./(repmat(Npost,1,size(bovar,2))-1))';
    bosem = (bovar'./repmat(Npost,1,size(bovar',2)))';
    postselection = (Npost./Nrealizations)';

    %automatic fit
    t_max = 38; % index maximum time used for fit
    disp(['Maximum time used for fit = ', num2str(tvalues(t_max)), 32, 'ms'])
    x = (1:size(bo(:,1:t_max),1))-centersite;
    y = tvalues(1:t_max)'; 
    x = repmat(x,1,length(y))';  
    y = repelem(y,length(x)/length(y))'; 
    z = reshape(bo(:,1:t_max),length(y),1);
    f = fit([x,y],z,'abs(besselj(abs(x),2*(2*a)/b*sin(pi*b*y/1000))).^2',...
        'StartPoint', [10, 22], 'Upper', [20, 30], 'Lower', [2, 10]); %original
    x = (1:size(bo,1))-centersite;
    y = tvalues';
    x = repmat(x,1,length(y))';  
    y = repelem(y,length(x)/length(y))'; 
    bo_fit = reshape(f([x,y]),size(bo,1), size(bo,2));
    JFit(vv,:) = [f.a f.b diff(confint(f,0.95))/2];
    

    %% Plot

    yaxis = linspace(-6,6,7);
    
    f1 = figure('name','2D2 quantum walk (0-73.5 ms)','Position',[400, 400, 300, 250]); 
    imagesc(tvalues(1:22), yaxis, bo(:,1:22), [0 1])
    yticks([-5 0 5])
    ylabel('Position i (sites)')
    axis square
    colorbar
   
    f2 = figure('name','2D2 quantum walk (125-177.5 ms)','Position',[400, 400, 300, 250]); 
    imagesc(tvalues(23:38), yaxis, bo(:,23:38), [0 1])
    yticks([-5 0 5])
    axis square
    colorbar
    
    f3 = figure('name','2D2 quantum walk (225-277.5 ms)','Position',[400, 400, 300, 250]); 
    imagesc(tvalues(39:54), yaxis, bo(:,39:54), [0 1])
    yticks([-5 0 5])
    ylabel('Position i (sites)')
    axis square
    colorbar

    f4 = figure('name','2D2 quantum walk (325-377.5 ms)','Position',[400, 400, 300, 250]); 
    imagesc(tvalues(55:70), yaxis, bo(:,55:70), [0 1])
    yticks([-5 0 5])
    axis square
    colorbar
    

    %%
    
    figure('name','2D2 quantum walk fit','units','normalized','position',[0.15,0.2,0.7,0.6]);
    subplot(2,2,1)
    imagesc(tvalues(1:22), row1:row2, bo_fit(:,1:22), [0 1])
    axis square
    xlabel('Time (ms)')
    ylabel('Row')
    colorbar
    J_str = ['J/h = ',num2str(f.a,'%.1f'),' Hz'];
    E_str = ['E/h = ',num2str(f.b,'%.1f'),' Hz'];
    title({J_str, E_str})

    subplot(2,2,2)
    imagesc(tvalues(23:38), row1:row2, bo_fit(:,23:38), [0 1])
    colorbar
    set(gca,'yticklabel',{[]})
    axis square
    xlabel('Time (ms)')

    subplot(2,2,3)
    imagesc(tvalues(39:54), row1:row2, bo_fit(:,39:54), [0 1])
    colorbar
    set(gca,'yticklabel',{[]})
    axis square
    xlabel('Time (ms)')

    subplot(2,2,4)
    imagesc(tvalues(55:70), row1:row2, bo_fit(:,55:70), [0 1])
    colorbar
    set(gca,'yticklabel',{[]})
    axis square
    xlabel('Time (ms)')
end


%% central density decay

central_density = zeros([2,70]);

central_density(1,:) = bo(7,:);
central_density(2,:) = bo_fit(7,:);

x = tvalues(1:38);
y = central_density(1,1:38);
fo = fitoptions('Method','NonlinearLeastSquares',...
               'Lower',[300],...
               'Upper',[500],...
               'StartPoint',[400]);
ft = fittype('0.4782*exp(-x/1000/(a/2000))*cos(2*pi*22.58*x/1000) + 0.4547','options',fo);
g = fit(x,y',ft); %fit the central density
tplot = linspace(0,400,1000);

fB = figure('Position',[400, 400, 300, 250]);
plot(tvalues(1:70), central_density(1,:),"o")
err = sqrt(bosem(7,:));
errorbar(tvalues(1:70), central_density(1,:),err,'.','MarkerSize',20,'CapSize',0)
hold on;
plot(tplot,g(tplot))
xlabel('Evolution time (ms)','Fontsize', 12)
ylabel('$\langle n_0 \rangle$','Fontsize', 12,'Interpreter', 'Latex')


%% Post-selection rate

plot_figure = 1;
if plot_figure
    figure()
    hold on
    plot(tvalues, postselection, '-o', 'Linewidth', 1.5)
    xlabel('Time (ms)')
    ylabel('Post-selection rate')
    ylim([0, 1])
    hold off
end
